package context

import (
	"errors"
	"fmt"
	"github.com/Unixeno/TheMoon/model"
	"github.com/Unixeno/TheMoon/moon/starlink"
	"sync"
	"sync/atomic"
)

var ClosedStream = errors.New("stream already closed")

type StreamCallback func(stream model.Stream) error

type streamContext struct {
	callbacks sync.Map // 记录流对应callback的map
	streams   sync.Map // 当前活动的流
}

func (s *streamContext) registerCallback(module string, operation int, callback StreamCallback) {
	s.callbacks.Store(s.genKey(module, operation), callback)
}

func (s *streamContext) registerStream(streamID string, stream *moonStream) {
	s.streams.Store(streamID, stream)
}

func (s *streamContext) unRegisterStream(streamID string) {
	s.streams.Delete(streamID)
}

func (s *streamContext) getStream(streamID string) *moonStream {
	stream, ok := s.streams.Load(streamID)
	if ok {
		return stream.(*moonStream)
	} else {
		return nil
	}
}

func (s *streamContext) getCallback(module string, operation int) StreamCallback {
	callback, ok := s.callbacks.Load(s.genKey(module, operation))
	if ok {
		return callback.(StreamCallback)
	} else {
		return nil
	}
}

func (s *streamContext) genKey(module string, operation int) string {
	return fmt.Sprintf("%s_%d", module, operation)
}

type moonStream struct {
	timeout int
	model.Message
	client    starlink.Client
	readChan  chan []byte
	handshake int32 // 0表示第一次发送
	closed    int32 // 1表示流已关闭
}

func newMoonStream(src string, module string, operation int, client starlink.Client) *moonStream {
	message := model.NewMessage(src, module)
	stream := &moonStream{
		client:   client,
		readChan: make(chan []byte, 4),
		Message:  message,
	}
	stream.Message.Operation = int32(operation)
	stream.Header.Stream = true
	stream.Header.Sync = false // 用于握手
	return stream
}

func getMoonStream(message *model.Message, client starlink.Client) *moonStream {
	stream := newMoonStream(message.Header.Destination, message.Header.Source, int(message.Operation), client)
	stream.handshake = 1 // 消息已经完成了握手
	stream.Header.Sync = true
	return stream
}

func (s *moonStream) ID() string {
	return s.Header.ID
}

func (s *moonStream) Send(msg []byte) error {
	// sync = 0, msg != nil 握手
	// sync = 1, msg != nil 消息
	// sync = 1, msg == nil 保持链接
	// sync = 0，msg == nil 关闭

	// 检查流是否关闭
	if atomic.LoadInt32(&s.closed) == 1 {
		if msg == nil {
			// 发送关闭消息
			logger.WithField("id", s.ID()).Debug("send close")
			s.Header.Sync = false
			s.Content = msg
			_ = s.client.Send(&s.Message)
			return nil
		} else {
			logger.WithField("id", s.ID()).Error("send to closed stream")
			return ClosedStream
		}
	}

	if 0 == atomic.SwapInt32(&s.handshake, 1) {
		// 第一次发送
		s.Content = msg
		if err := s.client.Send(&s.Message); err != nil {
			logger.WithField("id", s.ID()).WithError(err).Error("stream handshake failed")
			s.cleanup()
			return err
		}
		s.Header.Sync = true
	} else {
		s.Content = msg
		err := s.client.Send(&s.Message)
		if err != nil {
			// 发送失败，视为流断开
			s.cleanup()
			return err
		}
	}
	return nil
}

func (s *moonStream) write(msg []byte) error {
	// 检查流是否关闭
	if atomic.LoadInt32(&s.closed) == 1 {
		return ClosedStream
	}
	s.readChan <- msg
	return nil
}

func (s *moonStream) Recv() ([]byte, error) {
	msg, ok := <-s.readChan
	logger.Info("received")
	if ok {
		return msg, nil
	}
	return nil, ClosedStream
}

func (s *moonStream) cleanup() {
	if atomic.SwapInt32(&s.closed, 1) == 0 {
		logger.WithField("id", s.ID()).Debug("clean up stream")
		close(s.readChan)
		_ = s.Send(nil)
		localStreamContext.unRegisterStream(s.ID())
	}
}

func (s *moonStream) Close() error {
	s.cleanup()
	return nil
}
